package com.fengwk.cv4j.compute.ind;

import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.factory.Nd4j;
import org.nd4j.linalg.ops.transforms.Transforms;

import com.fengwk.cv4j.compute.Compute;
import com.fengwk.cv4j.utils.Nd4jUtils;
import com.fengwk.cv4j.view.ind.INDArrayWrap;
import com.fengwk.support.exception.BaseException;
import com.fengwk.support.function.Return;
import com.fengwk.support.util.ArrayUtils;
import com.fengwk.support.util.MathUtils;

/**
 * 支持INDArrayWrap的运算,应当支持INDArray和Number类型的各种结合方式
 * 
 * @author fengwk
 *
 */
public class INDCompute implements Compute<INDArrayWrap> {

	@Override
	public INDArrayWrap add(INDArrayWrap t1, INDArrayWrap t2) {
		INDArray arr1 = t1.arr(), arr2 = t2.arr();
		Number num1 = t1.num(), num2 = t2.num();
		if (arr1 == null && arr2 == null)
			return new INDArrayWrap(num1.doubleValue() + num2.doubleValue());
		if (arr1 == null && arr2 != null)
			return new INDArrayWrap(arr2.add(num1));
		if (arr1 != null && arr2 == null)
			return new INDArrayWrap(arr1.add(num2));
		else
			return sameShapeCompute(arr1, arr2, (a1, a2) -> new INDArrayWrap(a1.add(a2)));
	}

	@Override
	public INDArrayWrap sub(INDArrayWrap t1, INDArrayWrap t2) {
		INDArray arr1 = t1.arr(), arr2 = t2.arr();
		Number num1 = t1.num(), num2 = t2.num();
		if (arr1 == null && arr2 == null)
			return new INDArrayWrap(num1.doubleValue() - num2.doubleValue());
		if (arr1 == null && arr2 != null)
			return new INDArrayWrap(arr2.rsub(num1));
		if (arr1 != null && arr2 == null)
			return new INDArrayWrap(arr1.sub(num2));
		else
			return sameShapeCompute(arr1, arr2, (a1, a2) -> new INDArrayWrap(a1.sub(a2)));
	}

	@Override
	public INDArrayWrap mul(INDArrayWrap t1, INDArrayWrap t2) {
		INDArray arr1 = t1.arr(), arr2 = t2.arr();
		Number num1 = t1.num(), num2 = t2.num();
		if (arr1 == null && arr2 == null)
			return new INDArrayWrap(num1.doubleValue() * num2.doubleValue());
		if (arr1 == null && arr2 != null) {
			// 处理了0的情况为了性能的优化和转置矩阵的运算简洁
			if (num1.doubleValue() == 0)
				return new INDArrayWrap(0);
			return new INDArrayWrap(arr2.mul(num1));
		}
		if (arr1 != null && arr2 == null) {
			if (num2.doubleValue() == 0)
				return new INDArrayWrap(0);
			return new INDArrayWrap(arr1.mul(num2));
		}
		else
			return sameShapeCompute(arr1, arr2, (a1, a2) -> new INDArrayWrap(a1.mul(a2)));
	}

	@Override
	public INDArrayWrap div(INDArrayWrap t1, INDArrayWrap t2) {
		INDArray arr1 = t1.arr(), arr2 = t2.arr();
		Number num1 = t1.num(), num2 = t2.num();
		if (arr1 == null && arr2 == null)
			return new INDArrayWrap(num1.doubleValue() / num2.doubleValue());
		if (arr1 == null && arr2 != null)
			return new INDArrayWrap(arr2.rdiv(num1));
		if (arr1 != null && arr2 == null)
			return new INDArrayWrap(arr1.div(num2));
		else
			return sameShapeCompute(arr1, arr2, (a1, a2) -> new INDArrayWrap(a1.div(a2)));
	}

	@Override
	public INDArrayWrap pow(INDArrayWrap t1, INDArrayWrap t2) {
		INDArray arr1 = t1.arr(), arr2 = t2.arr();
		Number num1 = t1.num(), num2 = t2.num();
		if (arr1 == null && arr2 == null)
			return new INDArrayWrap(Math.pow(num1.doubleValue(), num2.doubleValue()));
		if (arr1 == null && arr2 != null) 
			return new INDArrayWrap(Transforms.pow(Nd4j.valueArrayOf(arr2.shape(), num1.doubleValue()), arr2));
		if (arr1 != null && arr2 == null)
			return new INDArrayWrap(Transforms.pow(arr1, num2));
		else
			return sameShapeCompute(arr1, arr2, (a1, a2) -> new INDArrayWrap(Transforms.pow(a1, a2)));
	}

	@Override
	public INDArrayWrap exp(INDArrayWrap t) {
		INDArray arr = t.arr();
		Number num = t.num();
		if (arr == null)
			return new INDArrayWrap(Math.exp(num.doubleValue()));
		else 
			return new INDArrayWrap(Transforms.exp(arr));
	}

	@Override
	public INDArrayWrap log(INDArrayWrap t) {
		INDArray arr = t.arr();
		Number num = t.num();
		if (arr == null)
			return new INDArrayWrap(Math.log(num.doubleValue()));
		else 
			return new INDArrayWrap(Transforms.log(arr));
	}

	@Override
	public INDArrayWrap neg(INDArrayWrap t) {
		INDArray arr = t.arr();
		Number num = t.num();
		if (arr == null)
			return new INDArrayWrap(- num.doubleValue());
		else 
			return new INDArrayWrap(Transforms.neg(arr));
	}
	
	@Override
	public INDArrayWrap mean(INDArrayWrap...ts) {
        if (ts == null || ts.length == 0)
            throw new BaseException("ts == null || ts.length == 0");
        
        INDArrayWrap wrap0 = ts[0];
        INDArray arr = wrap0.arr();
        if (arr == null) 
            return new INDArrayWrap(MathUtils.mean(toNumbers(ts)));
        else {
            if (ts.length > 1) 
            	return new INDArrayWrap(Nd4jUtils.mean(arr, toDimensionFrom1(ts)));
            else 
                return new INDArrayWrap(arr);
        }
    }
	
	@Override
	public INDArrayWrap sum(INDArrayWrap...ts) {
        if (ts == null || ts.length == 0)
            throw new BaseException("ts == null || ts.length == 0");
        
        INDArrayWrap wrap0 = ts[0];
        INDArray arr = wrap0.arr();
        if (arr == null) 
            return new INDArrayWrap(MathUtils.sum(toNumbers(ts)));
        else {
            if (ts.length > 1) 
            	return new INDArrayWrap(Nd4jUtils.sum(arr, toDimensionFrom1(ts)));
            else 
                return new INDArrayWrap(arr);
        }
    }

    @Override
    public INDArrayWrap mod(INDArrayWrap t1, INDArrayWrap t2) {
        INDArray arr1 = t1.arr(), arr2 = t2.arr();
        Number num1 = t1.num(), num2 = t2.num();
        if (arr1 == null && arr2 == null)
            return new INDArrayWrap(num1.doubleValue() % num2.doubleValue());
        if (arr1 == null && arr2 != null) 
            return new INDArrayWrap(Nd4j.valueArrayOf(arr2.shape(), num1.doubleValue()).fmod(arr2));
        if (arr1 != null && arr2 == null)
            return new INDArrayWrap(arr1.fmod(num2.doubleValue()));
        else
            return sameShapeCompute(arr1, arr2, (a1, a2) -> new INDArrayWrap(arr1.fmod(arr2)));
    }

    @Override
    public INDArrayWrap abs(INDArrayWrap t) {
        INDArray arr = t.arr();
        Number num = t.num();
        if (arr == null)
            return new INDArrayWrap(Math.abs(num.doubleValue()));
        else 
            return new INDArrayWrap(Transforms.abs(arr));
    }

    @Override
    public INDArrayWrap sign(INDArrayWrap t) {
        INDArray arr = t.arr();
        Number num = t.num();
        if (arr == null) 
            return new INDArrayWrap(MathUtils.sign(num));
        else 
            return new INDArrayWrap(Transforms.sign(arr));
    }

    @Override
    public INDArrayWrap max(INDArrayWrap...ts) {
        if (ts == null || ts.length == 0)
            throw new BaseException("ts == null || ts.length == 0");
        
        INDArrayWrap wrap0 = ts[0];
        INDArray arr = wrap0.arr();
        if (arr == null) 
            return new INDArrayWrap(MathUtils.maxNum(toNumbers(ts)));
        else {
            if (ts.length > 1) {
            	return new INDArrayWrap(Nd4jUtils.max(arr, toDimensionFrom1(ts)));
            } else 
                return new INDArrayWrap(arr);
        }
    }

    @Override
    public INDArrayWrap min(INDArrayWrap...ts) {
        if (ts == null || ts.length == 0)
            throw new BaseException("ts == null || ts.length == 0");
        
        INDArrayWrap wrap0 = ts[0];
        INDArray arr = wrap0.arr();
        if (arr == null) 
            return new INDArrayWrap(MathUtils.minNum(toNumbers(ts)));
        else {
            if (ts.length > 1) 
            	return new INDArrayWrap(Nd4jUtils.min(arr, toDimensionFrom1(ts)));
            else 
                return new INDArrayWrap(arr);
        }
    }

    @Override
    public INDArrayWrap cos(INDArrayWrap t) {
        INDArray arr = t.arr();
        Number num = t.num();
        if (arr == null) 
            return new INDArrayWrap(Math.cos(num.doubleValue()));
        else 
            return new INDArrayWrap(Transforms.cos(arr));
    }

    @Override
    public INDArrayWrap sin(INDArrayWrap t) {
        INDArray arr = t.arr();
        Number num = t.num();
        if (arr == null) 
            return new INDArrayWrap(Math.sin(num.doubleValue()));
        else 
            return new INDArrayWrap(Transforms.sin(arr));
    }

    @Override
    public INDArrayWrap tan(INDArrayWrap t) {
        INDArray arr = t.arr();
        Number num = t.num();
        if (arr == null) 
            return new INDArrayWrap(Math.tan(num.doubleValue()));
        else {
            arr = arr.dup();
            for (int i = 0; i < arr.length(); i ++)
                arr.putScalar(i, Math.tan(arr.getDouble(i)));
            return new INDArrayWrap(arr);
        }
    }

    @Override
    public INDArrayWrap atan(INDArrayWrap t) {
        INDArray arr = t.arr();
        Number num = t.num();
        if (arr == null) 
            return new INDArrayWrap(Math.atan(num.doubleValue()));
        else 
            return new INDArrayWrap(Transforms.atan(arr));
    }

    @Override
    public INDArrayWrap round(INDArrayWrap t) {
        INDArray arr = t.arr();
        Number num = t.num();
        if (arr == null) 
            return new INDArrayWrap(Math.round(num.doubleValue()));
        else 
            return new INDArrayWrap(Transforms.round(arr));
    }

    @Override
    public INDArrayWrap floor(INDArrayWrap t) {
        INDArray arr = t.arr();
        Number num = t.num();
        if (arr == null) 
            return new INDArrayWrap(Math.floor(num.doubleValue()));
        else 
            return new INDArrayWrap(Transforms.floor(arr));
    }

    @Override
    public INDArrayWrap ceil(INDArrayWrap t) {
        INDArray arr = t.arr();
        Number num = t.num();
        if (arr == null) 
            return new INDArrayWrap(Math.ceil(num.doubleValue()));
        else 
            return new INDArrayWrap(Transforms.ceil(arr));
    }
	
    @Override
    public INDArrayWrap equals(INDArrayWrap t1, INDArrayWrap t2) {
        INDArray arr1 = t1.arr(), arr2 = t2.arr();
        Number num1 = t1.num(), num2 = t2.num();
        INDArray buf = null;
        if (arr1 == null && arr2 == null)
            return new INDArrayWrap(num1.doubleValue() == num2.doubleValue() ? 1 : 0);
        if (arr1 == null && arr2 != null) {
        	buf = arr2.dup();
        	double d = num1.doubleValue();
        	for (int i = 0; i < buf.length(); i ++)
        		buf.putScalar(i, buf.getDouble(i) == d ? 1 : 0);
        }
        if (arr1 != null && arr2 == null) {
        	buf = arr1.dup();
        	double d = num2.doubleValue();
        	for (int i = 0; i < buf.length(); i ++)
        		buf.putScalar(i, buf.getDouble(i) == d ? 1 : 0);
        } else {
        	if (isSameShape(arr1, arr2) != -1)
        		throw new BaseException("isSameShape(arr1, arr2) != -1");
        	buf = arr1.dup();
        	for (int i = 0; i < buf.length(); i ++) {
                buf.putScalar(i, buf.getDouble(i) == arr2.getDouble(i) ? 1 : 0);
            }
        }
        return new INDArrayWrap(buf);
    }
    
    /* INDArray only */
    
    public INDArrayWrap mmul(INDArrayWrap t1, INDArrayWrap t2) {
        INDArray arr1 = t1.arr(), arr2 = t2.arr();
        if (arr1 == null || arr2 == null)
            throw new BaseException("arr1 == null || arr2 == null");
        return new INDArrayWrap(arr1.mmul(arr2));
    }
    
    /**
     * 将INDArrayWrap数组转为Number数组
     * 
     * @param ts
     * @return
     */
    private Number[] toNumbers(INDArrayWrap...ts) {
        if (ArrayUtils.isBlank(ts))
            throw new BaseException("ArrayUtils.isBlank(ts)");
        
        Number[] nums = new Number[ts.length];
        for (int i = 0; i < ts.length; i ++) {
            Number num = ts[i].num();
            if (num == null)
                throw new BaseException("num == null");
            nums[i] = num;
        }
        return nums;
    }
    
    /**
     * 将INDArrayWrap数组从第1号下标开始转为dimension数组
     * 
     * @param ts
     * @return
     */
    private int[] toDimensionFrom1(INDArrayWrap...ts) {
        if (ts == null || ts.length < 2)
            throw new BaseException("ts == null || ts.length < 2");
        
        int[] dimension = new int[ts.length - 1];
        for (int i = 1; i < ts.length; i ++) {
            Number num = ts[i].num();
            if (num == null)
                throw new BaseException("num == null");
            dimension[i - 1] = num.intValue();
        }
        return dimension;
    }
    
    /**
     * 使两个INDArray形状一致再进行运算
     * 在Nd4j中有一些奇怪的现象,例如下边两个形状不一致的矩阵运算,Nd4j好像会在第三列塞入一些随机数字去运算
     * {                            {
     *      {1/3f, 1/3f, 1/3f},         {0},
     *      {1/3f, 1/3f, 1/3f}          {0}
     * }                            }
     * 相当于在这里采用自动广播的形式进行运算
     * 
     * @param arr1
     * @param arr2
     * @return
     */
    private INDArrayWrap sameShapeCompute(INDArray arr1, INDArray arr2, Return.Params2<INDArrayWrap, INDArray, INDArray> computeArr) {
        int flag = isSameShape(arr1, arr2);
        if (flag == 0) 
            // arr1形状大于arr2
            arr2 = arr2.broadcast(arr1.shape());
        if (flag == 1)
            arr1 = arr1.broadcast(arr2.shape());
        // 运算函数
        return computeArr.apply(arr1, arr2);
    }
    
    private int isSameShape(INDArray arr1, INDArray arr2) {
        if (arr1 == null || arr2 == null)
            throw new BaseException("arr1 == null || arr2 == null");
        int[] shape1 = arr1.shape(), shape2 = arr2.shape();
        if (shape1.length != shape2.length)
            throw new BaseException("shape1.length != shape2.length");
        int flag = -1;// 0:1>2,1:1<2
        for (int i = 0; i < shape1.length; i ++) {
            int s1 = shape1[i], s2 = shape2[i];
            if (s1 == s2) 
                continue;
            else if (s1 > s2) {
                if (flag == 1)
                    throw new BaseException("s1 > s2 && flag == 1");
                flag = 0;
            } else if (s1 < s2) {
                if (flag == 0)
                    throw new BaseException("s1 < s2 && flag == 0");
                flag = 1;
            }
        }
        return flag;
    }

}
